package com.liuqi.openai.keyword;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Stack;
import java.util.stream.Collectors;

/**
 * 关键词提取器的抽象类
 *
 * 定义关键词存储器 {@link TrieNode}
 * 定义关键词提取结果 {@link MatchResult}
 * 定义模糊匹配状态 {@link FuzzyMatchState}
 *
 * 关键词中如果包含的字母, 会统一转成小写进行存储.
 *
 * 在检索文本提取关键词时, 也需要将字母转换成小写.
 *
 * 关键词匹配到结果, 返回词典中定义的关键词.
 *
 * @author liuqi
 * @date 2025/8/3
 **/
public abstract class AbstractKeywordExtractor implements KeywordExtractor {

    /**
     * 模糊匹配最大编辑距离
     */
    protected final static int FUZZY_MATCH_MAX_EDIT_DISTANCE = 3;

    /**
     * 关键词存储根节点
     */
    public final TrieNode root;

    /**
     * 已停用关键词
     */
    protected final Set<String> stopWords;

    /**
     * 关键词与字典映射关系
     */
    protected final Map<String, Dictionary> dictionaryMap;

    protected AbstractKeywordExtractor(List<Dictionary> dictionaries) {
        this(dictionaries, Collections.emptySet());
    }

    public AbstractKeywordExtractor(List<Dictionary> dictionaries, Set<String> stopWords) {
        if (dictionaries == null || dictionaries.isEmpty()) {
            throw new NullPointerException("keyword dictionaries cannot be empty");
        }
        this.root = new TrieNode(null, null);
        if (stopWords == null) {
            this.stopWords = Collections.emptySet();
        } else {
            this.stopWords = new HashSet<>();
            stopWords.forEach(word -> this.stopWords.add(word.toLowerCase()));
        }
        this.dictionaryMap = new HashMap<>();
        this.insertWords(dictionaries);
    }

    @Override
    public List<Dictionary> extract(String text) {
        if (text == null || text.isEmpty()) {
            return Collections.emptyList();
        }

        // 关键词提取器结果
        List<Dictionary> result = new ArrayList<>(2);

        int length = text.length();
        for (int i = 0; i < length; ) {
            // 空白字符不参与提取关键词
            if (Character.isWhitespace(text.charAt(i))) {
                i++;
                continue;
            }

            // 由子类进行匹配
            MatchResult match = match(text, i);
            if (match != null) {
                // 添加关键词提取结果
                addMatchResult(result, match);

                // 当前匹配索引加上已匹配完的字符长度
                i += match.length;
            } else {
                i++;
            }
        }

        return result;
    }

    /**
     * 批量添加关键词词到关键词存储树
     *
     * @param dictionaries 关键词词典
     */
    public void insertWords(List<Dictionary> dictionaries) {
        if (dictionaries != null) {
            dictionaries.forEach(this::insertWord);
        }
    }

    /**
     * 添加关键词词到关键词存储树
     *
     * @param dictionary 关键词词典
     */
    public void insertWord(Dictionary dictionary) {
        if (dictionary == null
                || dictionary.getKeyword() == null
                || dictionary.getKeyword().isEmpty()) {
            return;
        }
        // 关键词去掉空格
        String word = dictionary.getKeyword().trim();
        // 添加到Trie树
        TrieNode current = root;
        for (int i = 0; i < word.length(); i++) {
            // 当前节点树
            final TrieNode tmpCurrent = current;
            // 字母由小写进行存储
            char key = Character.toLowerCase(word.charAt(i));
            // 如果key不存在, 则进行添加
            current = current.children.computeIfAbsent(key, c -> new TrieNode(tmpCurrent, key));
        }
        // 标记最后一个key为结束状态
        current.setEndOfWord(true);

        // 存储关键词与字典映射, 关键词会转换为小写
        dictionaryMap.put(dictionary.getKeyword().toLowerCase(), dictionary);

        // 关键词为停用状态
        if (dictionary.isStop()) {
            stopWords.add(dictionary.getKeyword().toLowerCase());
        }
    }

    /**
     * 添加停用词
     *
     * @param word 关键词
     */
    public void insertStopWord(String word) {
        if (word != null && !word.isEmpty()) {
            stopWords.add(word.toLowerCase());
        }
    }

    /**
     * 检查是否是停用词
     */
    public boolean isStopWord(String word) {
        return stopWords.contains(word.toLowerCase());
    }

    /**
     * 从词典中移除词
     *
     * @param word 要移除的词
     */
    public void removeWord(String word) {
        if (word == null || word.isEmpty()) {
            return;
        }
        word = word.trim();
        removeWordFromTrie(word);
        this.dictionaryMap.remove(word);
    }

    /**
     * 由子类进行提取关键词
     *
     * @param text 文本
     * @param pos  文本起始位置
     * @return
     */
    protected abstract MatchResult match(String text, int pos);

    /**
     * 从关键词存储中移除关键词.
     */
    private void removeWordFromTrie(String word) {
        TrieNode current = root;
        Stack<TrieNode> nodeStack = new Stack<>();
        Stack<Character> charStack = new Stack<>();
        // 找到对应的节点路径
        for (int i = 0; i < word.length(); i++) {
            char ch = Character.toLowerCase(word.charAt(i));
            current = current.children.get(ch);
            if (current == null) {
                return; // 词不存在
            }
            nodeStack.push(current);
            charStack.push(ch);
        }

        // 对删除的节点先进行标记非结尾关键词状态
        current.setEndOfWord(false);

        // 如果节点没有子节点, 可以删除, 为关键词结尾
        if (current.endOfWord) {
            TrieNode parent;
            char ch;

            while (!nodeStack.isEmpty() && current.children.isEmpty()) {
                ch = charStack.pop();
                parent = nodeStack.isEmpty() ? root : nodeStack.pop().parent;
                parent.children.remove(ch);

                // 前置有重复的关键词
                // 例如: java、jav 当删除java关键词, 需要保留 jav 关键词
                if (parent.endOfWord) {
                    return;
                }
                current = parent;
            }


        }
    }

    /**
     * 将关键词匹配结果<code>match</code>构建成{@link DictionaryProxy}添加到结果集<code>result</code>
     *
     * @param result    结果集
     * @param match     关键词匹配结果
     */
    private void addMatchResult(final List<Dictionary> result, MatchResult match) {
        List<String> keywords = match.keywords.stream()
                .filter(word -> !isStopWord(word.toLowerCase()))
                .collect(Collectors.toList());

        if (keywords == null || keywords.isEmpty()) {
            return;
        }

        keywords.forEach(keyword -> {
            // 获取关键词词典
            Dictionary dict = dictionaryMap.get(keyword.toLowerCase());
            result.add(
                    DictionaryProxy.builder()
                            .keyword(dict.getKeyword())
                            .posTag(dict.getPosTag())
                            .related(dict.getRelated())
                            .keywords(match.keywords)
                            .matchLength(match.length)
                            .searchCount(match.searchCount)
                            .fuzzy(match.fuzzy)
                            .build()
            );
        });
    }

    /**
     * 关键词存储器, 是一个树结构
     * 将关键词进行拆分成单个字符, 递归存储到下游节点. 在进行关键词匹配时, 需要进行递归匹配到最后一个字符
     */
    protected class TrieNode {
        /**
         * 子节点
         * key: TrieNode
         */
        protected final Map<Character, TrieNode> children;

        /**
         * 父节点
         */
        protected final TrieNode parent;

        /**
         * 关键词细分的关键字
         */
        protected final Character key;

        /**
         * 关键词结尾标记
         */
        protected boolean endOfWord;

        TrieNode(TrieNode parent, Character key) {
            this.parent = parent;
            this.children = new HashMap<>();
            this.key = key;
        }

        void setEndOfWord(boolean endOfWord) {
            this.endOfWord = endOfWord;
        }
    }

    /**
     * 关键词提取结果
     */
    protected class MatchResult {
        /**
         * 匹配结果可能或提取到多个关键词
         */
        final List<String> keywords;

        /**
         * 文本匹配长度
         */
        final int length;

        /**
         * 提取到关键词的检索次数
         */
        final int searchCount;

        /**
         * 是否模仿匹配
         */
        final boolean fuzzy;

        public MatchResult(List<String> keywords, int length, int searchCount) {
            this(keywords, length, searchCount, false);
        }

        public MatchResult(List<String> keywords, int length, int searchCount, boolean fuzzy) {
            this.keywords = keywords;
            this.length = length;
            this.searchCount = searchCount;
            this.fuzzy = fuzzy;
        }
    }

    /**
     * 模糊匹配状态
     */
    protected class FuzzyMatchState {
        /**
         * 匹配节点
         */
        final TrieNode node;

        /**
         * 搜索偏移量
         */
        final int offset;

        /**
         * 当前关键词
         */
        final String currentWord;

        /**
         * 编辑距离
         */
        final int edits;

        FuzzyMatchState(TrieNode node, int offset, String currentWord, int edits) {
            this.node = node;
            this.offset = offset;
            this.currentWord = currentWord;
            this.edits = edits;
        }
    }

}
